# scripts/download_models.py
import os

from langchain_huggingface import HuggingFaceEmbeddings
import time


def download_embedding_model():
    """预下载嵌入模型"""
    print("开始下载嵌入模型...")
    project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    model_dir = os.path.join(project_root, "model_cache")
    start_time = time.time()
    os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
    try:
        # 下载模型
        embeddings = HuggingFaceEmbeddings(
            model_name="BAAI/bge-small-zh-v1.5",
            cache_folder=model_dir

        )

        # 测试模型
        test_embedding = embeddings.embed_query("测试句子")
        print(f"模型维度: {len(test_embedding)}")

        download_time = time.time() - start_time
        print(f"模型下载完成，耗时: {download_time:.2f}秒")

    except Exception as e:
        print(f"模型下载失败: {e}")
        raise


if __name__ == "__main__":
    download_embedding_model()
